from abc import ABC, abstractmethod
from typing import Iterator, Optional, TypeVar
from pydantic import BaseModel


BlockT = TypeVar('BlockT', bound='BaseBlock')


class BaseBlock(ABC):
    @classmethod
    def from_state(cls, data: dict) -> BlockT:
        """Get the block data from the state data.

        Returns:
            The block.

        """
        raise NotImplementedError()

    @abstractmethod
    async def run(self) -> Iterator[dict]:
        pass

    @abstractmethod
    def to_state(self, data: dict):
        """Update the state data based on the block data.

        Args:
            State data dictionary. IT WILL BE MUTATED.

        """
        ...


class ProgressChunk(BaseModel):
    done: bool = False
    total_steps: Optional[int] = None
    current_step: Optional[int] = None
    data: Optional[dict] = None
